package com.ruoyi.chat.sse;

import com.alibaba.fastjson2.JSONObject;
import com.ruoyi.chat.constant.RagFlowCode;
import com.ruoyi.common.utils.StringUtils;
import okhttp3.*;
import okio.BufferedSource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import javax.validation.constraints.NotNull;
import java.io.IOException;
import java.util.Objects;
import java.util.concurrent.TimeUnit;

@Component
public class OKHttpStream {

    private static final Logger log = LoggerFactory.getLogger(OKHttpStream.class);

    // http客户端连接时长 - 5分钟
    private static final Integer TIME_OUT = 5 * 60;

    // sseEmitter关闭标识 对应rag Flow-api返回data:true
    private static final String END_FLAG = "true";


    // 处理流式http接口-post请求
    public void chatStream(String streamChatUrl, SseEmitter sseEmitter, JSONObject params, String apiKey) {
        // 1. 创建 OkHttpClient 对象，并设置超时时间
        // todo 后续优化连接池
        OkHttpClient client = new OkHttpClient.Builder()
                .connectTimeout(TIME_OUT, TimeUnit.SECONDS)
                .readTimeout(TIME_OUT, TimeUnit.SECONDS)
                .writeTimeout(TIME_OUT, TimeUnit.SECONDS)
                .addInterceptor(new RetryInterceptor(3))
                .build();

        //  封装请求体参数
        RequestBody requestBody = RequestBody.create(Objects.requireNonNull(MediaType.parse("application/json; charset=utf-8")), params.toJSONString());

        // 封装请求头
        Headers headers = new Headers.Builder()
                .set("Content-Type","application/json")
                .set("Accept","text/event-stream")
                .set("Authorization",apiKey)
                .build();

        // 2. 构建 Request 对象
        Request request = new Request.Builder()
                .url(streamChatUrl)
                .headers(headers)
                .post(requestBody)
                .build();

        // 3. 创建 Call 对象
        Call call = client.newCall(request);

        // 4. 监听回调
        call.enqueue(new Callback() {
            @Override
            public void onFailure(@NotNull Call call, @NotNull IOException e) {
                log.error("进入 onFailure方法 {}", e.getMessage(), e);
                sseEmitter.completeWithError(e);
                // todo 抛出异常 全局捕获
            }

            @Override
            public void onResponse(@NotNull Call call, @NotNull Response response) {
                if (response.isSuccessful()) {
                    String chunkMessage;
                    try (ResponseBody responseBody = response.body()) {
                        BufferedSource source = null;
                        if (responseBody != null) {
                            source = responseBody.source();
                        }
                        while (source != null && !source.exhausted()) {
                            chunkMessage = source.readUtf8Line();
                            if (StringUtils.isBlank(chunkMessage)) {
                                continue;
                            }
                            // 此处需要截取sse导致的返回体前的data:，共5个字符
                            int index = chunkMessage.indexOf("{");
                            JSONObject jsonObject = JSONObject.parseObject(chunkMessage.substring(index));
                            if (null != jsonObject && null != jsonObject.get("code")) {
                                // rag flow返回码如果不是200 则直接停止
                                if (!StringUtils.equals(RagFlowCode.CODE_SUCCESS, jsonObject.getString("code"))) {
                                    log.error("ragFlow请求失败, 错误码： {}", jsonObject.getString("code"));
                                    sseEmitter.complete();
                                    break;
                                }
                            }
                            if (null != jsonObject && null != jsonObject.get("data")) {
                                // 异步发送给前端 防止阻塞
                                SseClient.send(sseEmitter, String.valueOf(jsonObject), apiKey);
                                // 如果data:true，表示rag flow推流结束 则关闭sse
                                if (StringUtils.equals(END_FLAG, jsonObject.getString("data"))) {
                                    sseEmitter.complete();
                                    break;
                                }
                            }

                        }
                    } catch (Exception e) {
                        log.error("解析失败： {}", e.getMessage());
                        sseEmitter.completeWithError(e);
                        // todo 抛出异常 全局捕获
                    }
                } else {
                    log.error("onResponse 方法请求失败 {}", response.message());
                    // todo 抛出异常 全局捕获
                }
            }
        });
    }


    /**
     * ragFlow重连拦截器
      */
    static class RetryInterceptor implements Interceptor{
        private final int maxRetries;
        private int retryCount = 0;

        public RetryInterceptor(int maxRetries) {
            this.maxRetries = maxRetries;
        }

        @NotNull
        @Override
        public Response intercept(@NotNull Chain chain) throws IOException {
            // todo 待调试 关闭ragflow的情况
            Request request = chain.request();
            Response response = chain.proceed(request);
            while (!response.isSuccessful() && retryCount < maxRetries){
                retryCount++;
                response = chain.proceed(request);
            }
            return response;
        }
    }
}
